from rag.rag_embeddings import BaseEmbeddings, OpenAIEmbeddings, LocalEmbeddings
import json

class RAGDatabase:
    def __init__(self, embeddings: BaseEmbeddings, use_cache=True):
        self.embeddings = embeddings
        self.documents = []
        if use_cache:
            self.load("C:\JaredLyu\Project\MiniMind\database\\rag_database.json")

    def load(self, file_path):
        with open(file_path, 'r', encoding='utf-8') as f:
            self.documents = json.load(f)

    def add_document(self, text):
        embedding = self.embeddings.get_embedding(text)
        self.documents.append({'text': text, 'embedding': embedding.cpu().numpy().tolist()})
    
    def persist(self, file_path):
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(self.documents, f, ensure_ascii=False, indent=4)

    def search(self, query, top_k=5):
        query_embedding = self.embeddings.get_embedding(query)
        similarities = [(doc['text'], self.embeddings.cal_similarity(doc['embedding'], query_embedding)) for doc in self.documents]
        similarities.sort(key=lambda x: x[1], reverse=True)
        return similarities[:top_k]